#pragma once

#include <windows.h>

enum class PeArch {
  Invalid,
  Native86,
  Native64,
  Dotnet86,
  Dotnet64,
  DotnetAnyCpu,
  DotnetAnyCpuPrefer32
};

// Thanks to blaquee for the investigative work!
static PeArch GetPeArch(const wchar_t* szFileName) {
  auto result = PeArch::Invalid;
  auto hFile = CreateFileW(szFileName, GENERIC_READ, FILE_SHARE_READ, nullptr,
                           OPEN_EXISTING, 0, nullptr);
  if (hFile != INVALID_HANDLE_VALUE) {
    // SEC_IMAGE will load the file like the Windows loader would, saving us RVA
    // -> File offset conversion crap.
    auto hMappedFile = CreateFileMappingW(
        hFile, nullptr, PAGE_READONLY | SEC_IMAGE, 0, 0, nullptr);
    if (hMappedFile) {
      auto fileMap = MapViewOfFile(hMappedFile, FILE_MAP_READ, 0, 0, 0);
      if (fileMap) {
        __try {
          auto pidh = PIMAGE_DOS_HEADER(fileMap);
          if (pidh->e_magic == IMAGE_DOS_SIGNATURE) {
            auto pnth = PIMAGE_NT_HEADERS(ULONG_PTR(fileMap) + pidh->e_lfanew);
            if (pnth->Signature == IMAGE_NT_SIGNATURE) {
              auto Machine = pnth->FileHeader.Machine;
              if (Machine == IMAGE_FILE_MACHINE_I386 ||
                  Machine == IMAGE_FILE_MACHINE_AMD64) {
                auto isDll = (pnth->FileHeader.Characteristics &
                              IMAGE_FILE_DLL) == IMAGE_FILE_DLL;
                auto isFile86 = Machine == IMAGE_FILE_MACHINE_I386;

                // Set the native architecture of the PE file (to still have
                // something to show for if the COM directory is invalid).
                result = isFile86 ? PeArch::Native86 : PeArch::Native64;

                // Get the address and size of the COM (.NET) directory.
                ULONG_PTR comAddr = 0, comSize = 0;
                if (isFile86)  // x86
                {
                  auto pnth32 = PIMAGE_NT_HEADERS32(pnth);
                  comAddr =
                      pnth32->OptionalHeader
                          .DataDirectory[IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR]
                          .VirtualAddress;
                  comSize =
                      pnth32->OptionalHeader
                          .DataDirectory[IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR]
                          .Size;
                } else  // x64
                {
                  auto pnth64 = PIMAGE_NT_HEADERS64(pnth);
                  comAddr =
                      pnth64->OptionalHeader
                          .DataDirectory[IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR]
                          .VirtualAddress;
                  comSize =
                      pnth64->OptionalHeader
                          .DataDirectory[IMAGE_DIRECTORY_ENTRY_COM_DESCRIPTOR]
                          .Size;
                }

                // Check if the file has a (valid) COM (.NET) directory.
                if (comAddr && comSize >= sizeof(IMAGE_COR20_HEADER)) {
                  // Find out which flavor of dotnet we're dealing with.
                  // Specifically, "Any CPU" can be compiled with two flavors,
                  // "Prefer 32 bit" or not. Without the 32bit preferred flag,
                  // the loader will load the .NET environment based on the
                  // current platforms bitness (x86 or x64) Class libraries
                  // (DLLs) cannot specify the "Prefer 32 bit".
                  // https://mega.nz/#!vx5nVILR!jLafWGWhhsC0Qo5fE-3oEIc-uHBcRpraOo8L_KlUeXI
                  // Binaries that do not have COMIMAGE_FLAGS_ILONLY appear to
                  // be executed in a process that matches their native type.
                  // https://github.com/x64dbg/x64dbg/issues/1758

                  auto pcorh =
                      PIMAGE_COR20_HEADER(ULONG_PTR(fileMap) + comAddr);
                  if (pcorh->cb == sizeof(IMAGE_COR20_HEADER)) {
                    auto flags = pcorh->Flags;
#define test(x) (flags & x) == x
#define MY_COMIMAGE_FLAGS_32BITPREFERRED 0x00020000
                    if (isFile86)  // x86
                    {
                      if (test(COMIMAGE_FLAGS_32BITREQUIRED))
                        result =
                            !isDll && test(MY_COMIMAGE_FLAGS_32BITPREFERRED)
                                ? PeArch::DotnetAnyCpuPrefer32
                                : PeArch::Dotnet86;
                      else if (test(COMIMAGE_FLAGS_ILONLY))
                        result = PeArch::DotnetAnyCpu;
                      else
                        result = PeArch::Dotnet86;
                    } else  // x64
                      result = PeArch::Dotnet64;
#undef MY_COMIMAGE_FLAGS_32BITPREFERRED
#undef test
                  }
                }
              }
            }
          }
        } __except (EXCEPTION_EXECUTE_HANDLER) {
        }
        UnmapViewOfFile(fileMap);
      }
      CloseHandle(hMappedFile);
    }
    CloseHandle(hFile);
  }
  return result;
}